{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensorflow 在训练期间使用checkpoint保存模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用程序检查点（checkpoint）技术，可以在训练过程中保存模型：\n",
    "* 训练程序崩溃，也不需要从头开始训练，加载一个检查点模型继续开始训练\n",
    "* 预估服务可以加载一个检查点模型实现模型更新"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. 读取数据构建模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 读取数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n",
      "11493376/11490434 [==============================] - 17s 1us/step\n"
     ]
    }
   ],
   "source": [
    "(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",
    "\n",
    "train_labels = train_labels[:10000]\n",
    "test_labels = test_labels[:10000]\n",
    "\n",
    "train_images = train_images[:10000].reshape(-1, 28 * 28) / 255.0\n",
    "test_images = test_images[:10000].reshape(-1, 28 * 28) / 255.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([5, 0, 4, 1, 9], dtype=uint8)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 这是个10分类的训练任务\n",
    "train_labels[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 定义简单模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义一个简单的序列模型\n",
    "def create_model():\n",
    "    model = tf.keras.models.Sequential([\n",
    "        keras.layers.Dense(512, activation='relu', input_shape=(784,)),\n",
    "        keras.layers.Dense(10, activation='softmax')\n",
    "    ])\n",
    "\n",
    "    model.compile(optimizer='rmsprop',\n",
    "                loss='sparse_categorical_crossentropy',\n",
    "                metrics=['accuracy'])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. 在训练期间保存模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "tf.keras.callbacks.ModelCheckpoint 允许在训练的过程中和结束时回调保存的模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-03-25 16:08:07.982856: I tensorflow/core/platform/cpu_feature_guard.cc:145] This TensorFlow binary is optimized with Intel(R) MKL-DNN to use the following CPU instructions in performance critical operations:  SSE4.1 SSE4.2 AVX AVX2 FMA\n",
      "To enable them in non-MKL-DNN operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2022-03-25 16:08:07.984969: I tensorflow/core/common_runtime/process_util.cc:115] Creating new thread pool with default inter op setting: 12. Tune using inter_op_parallelism_threads for best performance.\n"
     ]
    }
   ],
   "source": [
    "# 创建一个基本的模型实例\n",
    "model = create_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建一个保存模型权重的回调\n",
    "checkpoint_path = \"./traing_ckpt/cp_{epoch:02d}.ckpt\"\n",
    "cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,\n",
    "                                                 save_weights_only=True,\n",
    "                                                 verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 10000 samples, validate on 10000 samples\n",
      "Epoch 1/10\n",
      " 9888/10000 [============================>.] - ETA: 0s - loss: 0.4012 - accuracy: 0.8829\n",
      "Epoch 00001: saving model to ./traing_ckpt/cp_01.ckpt\n",
      "10000/10000 [==============================] - 4s 425us/sample - loss: 0.4007 - accuracy: 0.8830 - val_loss: 0.3100 - val_accuracy: 0.9074\n",
      "Epoch 2/10\n",
      " 9888/10000 [============================>.] - ETA: 0s - loss: 0.1827 - accuracy: 0.9447\n",
      "Epoch 00002: saving model to ./traing_ckpt/cp_02.ckpt\n",
      "10000/10000 [==============================] - 3s 265us/sample - loss: 0.1837 - accuracy: 0.9448 - val_loss: 0.1905 - val_accuracy: 0.9432\n",
      "Epoch 3/10\n",
      " 9952/10000 [============================>.] - ETA: 0s - loss: 0.1162 - accuracy: 0.9656\n",
      "Epoch 00003: saving model to ./traing_ckpt/cp_03.ckpt\n",
      "10000/10000 [==============================] - 2s 244us/sample - loss: 0.1156 - accuracy: 0.9658 - val_loss: 0.1702 - val_accuracy: 0.9497\n",
      "Epoch 4/10\n",
      " 9792/10000 [============================>.] - ETA: 0s - loss: 0.0819 - accuracy: 0.9759\n",
      "Epoch 00004: saving model to ./traing_ckpt/cp_04.ckpt\n",
      "10000/10000 [==============================] - 3s 255us/sample - loss: 0.0812 - accuracy: 0.9761 - val_loss: 0.1531 - val_accuracy: 0.9557\n",
      "Epoch 5/10\n",
      " 9856/10000 [============================>.] - ETA: 0s - loss: 0.0579 - accuracy: 0.9830\n",
      "Epoch 00005: saving model to ./traing_ckpt/cp_05.ckpt\n",
      "10000/10000 [==============================] - 3s 271us/sample - loss: 0.0578 - accuracy: 0.9831 - val_loss: 0.1615 - val_accuracy: 0.9537\n",
      "Epoch 6/10\n",
      " 9952/10000 [============================>.] - ETA: 0s - loss: 0.0409 - accuracy: 0.9892\n",
      "Epoch 00006: saving model to ./traing_ckpt/cp_06.ckpt\n",
      "10000/10000 [==============================] - 3s 270us/sample - loss: 0.0411 - accuracy: 0.9892 - val_loss: 0.1609 - val_accuracy: 0.9553\n",
      "Epoch 7/10\n",
      " 9824/10000 [============================>.] - ETA: 0s - loss: 0.0260 - accuracy: 0.9926\n",
      "Epoch 00007: saving model to ./traing_ckpt/cp_07.ckpt\n",
      "10000/10000 [==============================] - 3s 254us/sample - loss: 0.0260 - accuracy: 0.9925 - val_loss: 0.1799 - val_accuracy: 0.9548\n",
      "Epoch 8/10\n",
      " 9824/10000 [============================>.] - ETA: 0s - loss: 0.0182 - accuracy: 0.9945\n",
      "Epoch 00008: saving model to ./traing_ckpt/cp_08.ckpt\n",
      "10000/10000 [==============================] - 3s 266us/sample - loss: 0.0183 - accuracy: 0.9945 - val_loss: 0.1709 - val_accuracy: 0.9579\n",
      "Epoch 9/10\n",
      " 9920/10000 [============================>.] - ETA: 0s - loss: 0.0132 - accuracy: 0.9961\n",
      "Epoch 00009: saving model to ./traing_ckpt/cp_09.ckpt\n",
      "10000/10000 [==============================] - 3s 280us/sample - loss: 0.0132 - accuracy: 0.9961 - val_loss: 0.1906 - val_accuracy: 0.9601\n",
      "Epoch 10/10\n",
      " 9920/10000 [============================>.] - ETA: 0s - loss: 0.0089 - accuracy: 0.9976\n",
      "Epoch 00010: saving model to ./traing_ckpt/cp_10.ckpt\n",
      "10000/10000 [==============================] - 3s 284us/sample - loss: 0.0089 - accuracy: 0.9976 - val_loss: 0.1673 - val_accuracy: 0.9627\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7fa26c44f128>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 训练模型，训练过程中每个epoch会保存checkpoint\n",
    "model.fit(train_images, \n",
    "          train_labels,  \n",
    "          epochs=10,\n",
    "          validation_data=(test_images,test_labels),\n",
    "          callbacks=[cp_callback])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10000/1 - 1s - loss: 0.0840 - accuracy: 0.9627\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[0.1672822047203168, 0.9627]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.evaluate(test_images,  test_labels, verbose=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. 使用checkpoint文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建一个新model\n",
    "new_model = create_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 加载模型预估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7fa2804a5080>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 加载权重，不需要写.index等后缀\n",
    "new_model.load_weights(\"./traing_ckpt/cp_09.ckpt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10000/1 - 1s - loss: 0.0969 - accuracy: 0.9601\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[0.1905901687839796, 0.9601]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 直接进入评估\n",
    "new_model.evaluate(test_images,  test_labels, verbose=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[3.59074397e-13, 2.00585426e-18, 6.06315054e-10, 4.93331775e-08,\n",
       "        7.06807305e-18, 2.60786374e-13, 1.14687445e-20, 1.00000000e+00,\n",
       "        1.36654211e-13, 2.49361576e-09],\n",
       "       [5.73253167e-10, 8.55602877e-09, 1.00000000e+00, 2.37206201e-08,\n",
       "        8.17579501e-21, 1.47403212e-09, 1.38319738e-08, 6.44906981e-18,\n",
       "        6.63520738e-10, 4.05923955e-18],\n",
       "       [3.15703186e-10, 9.99622583e-01, 2.13169798e-04, 3.21664425e-06,\n",
       "        8.69879500e-07, 8.32690148e-08, 3.91072632e-07, 3.89748311e-05,\n",
       "        1.20672506e-04, 5.53059074e-08]], dtype=float32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 直接进行预估\n",
    "new_model.predict(test_images[:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 继续开始训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 10000 samples, validate on 10000 samples\n",
      "Epoch 1/3\n",
      "10000/10000 [==============================] - 3s 313us/sample - loss: 0.0090 - accuracy: 0.9977 - val_loss: 0.1743 - val_accuracy: 0.9609\n",
      "Epoch 2/3\n",
      "10000/10000 [==============================] - 3s 274us/sample - loss: 0.0048 - accuracy: 0.9988 - val_loss: 0.1875 - val_accuracy: 0.9578\n",
      "Epoch 3/3\n",
      "10000/10000 [==============================] - 3s 275us/sample - loss: 0.0040 - accuracy: 0.9991 - val_loss: 0.1947 - val_accuracy: 0.9591\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7fa25d98f160>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 使用新的回调训练模型，loss会接着当前状态继续训练\n",
    "new_model.fit(train_images, \n",
    "          train_labels,  \n",
    "          epochs=3,\n",
    "          validation_data=(test_images,test_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
